#include <windows.h>
#include <iostream>
#include "AMan_i.h"
#include <Wtsapi32.h>
#include <atlsafe.h>
#include <fstream>
#include <conio.h>

using namespace std;


HRESULT CreateInst(IManager **man)
{
	CComPtr<IManager> tmp;
	HRESULT hr = CoCreateInstance(__uuidof(Manager), NULL, CLSCTX_INPROC_SERVER, __uuidof(IManager), (void**)&tmp);
	if( hr != S_OK )
	{
			cout<<"Cannot create Instance("<<hr<<")\n";
			return S_FALSE;
	}
	*man = tmp.Detach();
	cout<<"Created OK\n";
	return S_OK;
}

HRESULT AddServer(BSTR sName, IManager *man)
{

	man->AddServer(sName);
	cout<<"Added OK\n";
	return S_OK;
}

HRESULT GetServer(IManager *man, IServer **svr)
{

	CComPtr<IServer> tmp;
	HRESULT res = man->GetServerByName(L"192.168.1.3", &tmp);

	*svr = tmp.Detach();
	cout<<"Get OK\n";
	return res;
}

HRESULT GetSessions(IServer *serv, SAFEARRAY ** ss)
{
	CComSafeArray<VARIANT> sarr;
	SAFEARRAY *sa;
	serv->GetSessions(&sa);
	sarr.Attach(sa);

	sarr.CopyTo(ss);
	cout<<"It's found "<<sarr.GetCount()<<" sessions\n";
	return S_OK;
}

HRESULT GetProcesses(ISessions *is, SAFEARRAY ** ss)
{
	CComSafeArray<VARIANT> parr;
	SAFEARRAY *pa;
	is->GetProcesses(&pa);
	parr.Attach(pa);

	parr.CopyTo(ss);
	cout<<"It's found "<<parr.GetCount()<<" processes\n";
	return S_OK;
}

void PrintProcesses(CComSafeArray<VARIANT>& pr)
{
	cout<<"PID\tNAME"<<endl;
	for(int i=0; i<pr.GetCount(); i++)
	{
		CComQIPtr<IProcesses> pptr = pr.GetAt(i).pdispVal;
		DWORD pid;
		BSTR name;
		pptr->GetProcessId(&pid);
		pptr->GetProcessName(&name);
		wcout<<i<<")"<<pid<<"\t"<<name<<endl;
	}
}

void KillProcess(IProcesses *pr)
{
	pr->Terminate();
}


void menu()
{
	int ch, ses;
	CComPtr<IManager> man;
	CComSafeArray<VARIANT> sarr;
	SAFEARRAY *sa;
	CComPtr<IServer> serv;
	CComQIPtr<ISessions> sptr;

	CComSafeArray<VARIANT> parr;
	SAFEARRAY *pa=0;

	while(1)
	{
	cout<<"0 - Create Manager Instance\n"
	<<"1 - Add server\n"
	<<"2 - Get Server Interface\n"
	<<"3 - Get Sessions\n"
	<<"4 - Get Specified Session\n"
	<<"5 - Get Info\n"
	<<"6 - Turn off session\n"
	<<"7 - Get Processes\n"
	<<"8 - Show Processes\n"
	<<"9 - Shut down server\n"
	<<"10 - Kill process\n"
	<<"Other -exit\n"; 
	cin>>ch;
	switch(ch)
	{
		case 0:
		{
			if( CreateInst(&man) == S_FALSE )
				return;
			break;
		}
		case 1:
		{
			AddServer(L"192.168.1.3", man);
			break;
		}
		case 2:
		{
			GetServer(man, &serv);
			break;
		}
		case 3:
		{
			GetSessions(serv, &sa);
			sarr.Attach(sa);
			break;
		}
		case 4:
		{
			cout<<"Select Session: ";
			cin>>ses;
			sptr = sarr.GetAt(ses).pdispVal;
			break;
		}
		case 5:
		{
				int sch;
				cout<<"1 - Station name info\n"
				<<"2 - User name info\n"
				<<"3 - Directory info\n"
				<<"4 - Protocol info\n"
				<<"5 - Address info\n"
				<<"6 - Display info\n"
				<<"Other - return\n";
				cin>>sch;
				switch(sch)
				{
					case 1:
					{
						BSTR sname;
						sptr->GetStationNameInfo(&sname);
						wcout<<"Station Name:"<<sname<<endl;
						break;
					}
					case 2:
					{
						BSTR uname;
						sptr->GetUserNameInfo(&uname);
						wcout<<"User Name:"<<uname<<endl;
						break;
					}
					case 3:
					{
						BSTR dname;
						sptr->GetDirectoryInfo(&dname);
						wcout<<"Dir Name:"<<dname<<endl;
						break;
					}
					case 4:
					{
						BSTR pr;
						sptr->GetProtocolInfo(&pr);
						wcout<<"Station Name:"<<pr<<endl;
						break;
					}
					case 5:
					{
						WTS_CLIENT_ADDRESS wa = {};
						sptr->GetAddressInfo(&wa);
						wcout<<"Address family:"<<wa.AddressFamily<<endl;
						break;
					}
					case 6:
					{
						WTS_CLIENT_DISPLAY wa = {};
						sptr->GetDisplayInfo(&wa);
						wcout<<"Color Depth:"<<wa.ColorDepth<<endl;
						wcout<<"Horizontal Resolution:"<<wa.HorizontalResolution<<endl;
						wcout<<"Vertical Resolution:"<<wa.VerticalResolution<<endl;
						break;
					}
					default:
						break;
				}
				getch();
		}
		case 6:
		{
			sptr->Logoff();
			break;
		}
		case 7:
		{
			GetProcesses(sptr, &pa);
			parr.Attach(pa);
			break;
		}
		case 8:
		{
			PrintProcesses(parr);
			break;
		}
		case 9:
		{
			serv->Shutdown();
			break;
		}
		case 10:
		{
			int id;
			cout<<"Select Index to kill: ";
			cin>>id;
			CComQIPtr<IProcesses> pkill = parr.GetAt(id).pdispVal;
			KillProcess(pkill);
			break;
		}
		default:
			return;
	}
	getch();
	system("cls");
	}
}


int main()
{
    CoInitialize(NULL);

	menu();

    CoUninitialize();

	return 0;
}